#!/usr/bin/env python3
"""
ClamAV病毒库全量下载与上传脚本
功能：
1. 使用wget下载全量最新病毒库
2. 上传病毒库文件到FTP服务器
3. 日志记录
4. 错误处理
"""

import os
import sys
import subprocess
import time
import logging
from datetime import datetime

# ===== 配置区域 =====
CLAMAV_DB_DIR = "/var/lib/clamav"  # 病毒库目录
BACKUP_DIR = "/opt/clamav_backups"  # 备份保存目录
LOG_FILE = "/var/log/clamav_updater.log"  # 日志文件路径
MAX_BACKUPS = 3  # 最大保留备份数量
# 病毒库下载配置
CLAMAV_DB_URL = "https://database.clamav.net"  # ClamAV病毒库官方下载地址
MAIN_CVD = "main.cvd"  # 主病毒库文件
DAILY_CVD = "daily.cvd"  # 每日更新病毒库文件
BYTECODE_CVD = "bytecode.cvd"  # 字节码病毒库文件
# 打包配置
PACKAGE_DIR = "/tmp/clamav_packages"  # 打包临时目录
PACKAGE_NAME_PREFIX = "clamav_db"  # 打包文件名前缀
# FTP服务器配置
FTP_HOST = '192.168.2.250'    # FTP服务器地址
FTP_USER = '51'        # FTP用户名
FTP_PASS = '51'        # FTP密码
FTP_PATH = '/版本发布区-Manual/clamav_db'  # FTP上传路径
FTP_ENCODING = 'gbk'        # FTP连接编码
# ===================

def setup_logging():
    """配置日志记录"""
    logging.basicConfig(
        filename=LOG_FILE,
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))

def run_command(cmd):
    """执行系统命令并处理错误"""
    try:
        result = subprocess.run(
            cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        logging.info(f"命令执行成功: {' '.join(cmd)}")
        logging.debug(f"输出: {result.stdout}")
        return True
    except subprocess.CalledProcessError as e:
        logging.error(f"命令执行失败: {' '.join(cmd)}")
        logging.error(f"错误代码: {e.returncode}")
        logging.error(f"错误输出: {e.stderr.strip()}")
        return False

def update_virus_db():
    """使用wget下载全量最新的ClamAV病毒库"""
    logging.info("开始下载全量最新病毒库...")
    
    # 确保病毒库目录存在
    if not os.path.exists(CLAMAV_DB_DIR):
        os.makedirs(CLAMAV_DB_DIR, exist_ok=True)
        logging.info(f"创建病毒库目录: {CLAMAV_DB_DIR}")
    
    # 要下载的病毒库文件列表
    cvd_files = [MAIN_CVD, DAILY_CVD, BYTECODE_CVD]
    download_success = True
    
    for cvd_file in cvd_files:
        file_url = f"{CLAMAV_DB_URL}/{cvd_file}"
        file_path = os.path.join(CLAMAV_DB_DIR, cvd_file)
        
        logging.info(f"正在下载: {file_url}")
        
        # 使用wget下载文件，添加用户代理头避免403错误
        cmd = ["wget", "-O", file_path, "-U", "clamav/1.0", file_url]
        if not run_command(cmd):
            logging.error(f"下载 {cvd_file} 失败")
            download_success = False
            continue
        
        # 检查文件是否下载成功
        if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
            logging.info(f"成功下载 {cvd_file}")
        else:
            logging.error(f"{cvd_file} 下载失败或文件为空")
            download_success = False
    
    if download_success:
        logging.info("病毒库下载完成")
        return True
    else:
        logging.error("病毒库下载失败")
        return False

def get_cvd_files():
    """获取病毒库目录中的cvd文件列表"""
    cvd_files = []
    for db_file in os.listdir(CLAMAV_DB_DIR):
        if db_file.endswith('.cvd'):
            cvd_files.append(db_file)
    return cvd_files

def upload_to_ftp(file_path):
    """上传文件到FTP服务器"""
    
    from ftplib import FTP
    import os
    
    logging.info(f"开始上传文件到FTP服务器: {file_path}")
    
    # 创建FTP连接
    ftp = FTP()
    ftp.encoding = FTP_ENCODING
    ftp.set_debuglevel(0)
    ftp.connect(FTP_HOST)
    ftp.login(FTP_USER, FTP_PASS)    
    # 上传文件
    filename = os.path.basename(file_path)
 
    ftp.cwd(FTP_PATH)    
    with open(file_path, 'rb') as f:
        ftp.storbinary(f'STOR  {filename}', f)
    
    ftp.quit()
    logging.info("文件上传成功")
    return True

def package_virus_db():
    """打包病毒库文件为压缩包"""
    import zipfile
    from datetime import datetime
    
    logging.info("开始打包病毒库文件...")
    
    # 确保打包目录存在
    if not os.path.exists(PACKAGE_DIR):
        os.makedirs(PACKAGE_DIR, exist_ok=True)
        logging.info(f"创建打包目录: {PACKAGE_DIR}")
    
    # 生成打包文件名（包含日期时间）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    package_name = f"{PACKAGE_NAME_PREFIX}_{timestamp}.zip"
    package_path = os.path.join(PACKAGE_DIR, package_name)
    
    try:
        # 创建zip压缩包
        with zipfile.ZipFile(package_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            # 添加所有cvd文件到压缩包
            cvd_files = get_cvd_files()
            for cvd_file in cvd_files:
                file_path = os.path.join(CLAMAV_DB_DIR, cvd_file)
                if os.path.exists(file_path):
                    zipf.write(file_path, cvd_file)
                    logging.info(f"添加文件到压缩包: {cvd_file}")
        
        logging.info(f"病毒库打包完成: {package_path}")
        return package_path
    
    except Exception as e:
        logging.error(f"打包病毒库失败: {str(e)}")
        return None

def cleanup_old_packages():
    """清理旧的打包文件"""
    import glob
    from datetime import datetime, timedelta
    
    logging.info("开始清理旧的打包文件...")
    
    if not os.path.exists(PACKAGE_DIR):
        return
    
    # 获取所有打包文件
    package_pattern = os.path.join(PACKAGE_DIR, f"{PACKAGE_NAME_PREFIX}_*.zip")
    package_files = glob.glob(package_pattern)
    
    # 按修改时间排序
    package_files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
    
    # 保留最新的MAX_BACKUPS个文件，删除其余的
    for i, package_file in enumerate(package_files):
        if i >= MAX_BACKUPS:
            try:
                os.remove(package_file)
                logging.info(f"删除旧打包文件: {package_file}")
            except Exception as e:
                logging.error(f"删除文件失败 {package_file}: {str(e)}")
    
    logging.info("清理完成")
        

def main():
    setup_logging()
    logging.info("=" * 50)
    logging.info("开始ClamAV病毒库更新任务")
    
    # 更新病毒库
    if not update_virus_db():
        logging.error("更新任务中止")
        return
    
    # 打包病毒库
    package_path = package_virus_db()
    if not package_path:
        logging.error("打包失败，任务中止")
        return
    
    # 上传打包文件到FTP
    if upload_to_ftp(package_path):
        logging.info("打包文件上传成功")
    else:
        logging.error("打包文件上传失败")
    
    # 清理旧的打包文件
    cleanup_old_packages()
    
    logging.info("任务完成\n")

if __name__ == "__main__":
    main()